import torch
import torch.nn as nn

# 残差块：2D卷积1+2D卷积2+残差连接
class ResBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ResBlock, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels, out_channels, 3, 1, 1), # 2D卷积：卷积核3*3，步长1，填充1
            torch.nn.BatchNorm2d(out_channels), # 批归一化：输入输出通道32
            torch.nn.ReLU() # ReLU激活函数
        )
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(out_channels, out_channels, 3, 1, 1), # 2D卷积：卷积核3*3，步长1，填充1
            torch.nn.BatchNorm2d(out_channels), # 批归一化：输入输出通道32
        )
        self.relu = torch.nn.ReLU()

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.conv2(out)
        out += residual
        out = self.relu(out)
        return out

# 3D卷积块：3D卷积*3
class ThreeDBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ThreeDBlock, self).__init__()
        self.conv20 = torch.nn.Sequential( # 尺寸/2
            torch.nn.Conv3d(in_channels, out_channels, 3, 2, 1), # 3D卷积：卷积核3*3*3，步长2，填充1
            torch.nn.BatchNorm3d(out_channels), # 批归一化
            torch.nn.ReLU() # ReLU激活函数
        )

        self.conv21 = torch.nn.Sequential(
            torch.nn.Conv3d(out_channels, out_channels, 3, 1, 1), # 3D卷积：卷积核3*3*3，步长1，填充1
            torch.nn.BatchNorm3d(out_channels), # 批归一化
            torch.nn.ReLU() # ReLU激活函数
        )

        self.conv22 = torch.nn.Sequential(
            torch.nn.Conv3d(out_channels, out_channels, 3, 1, 1), # 3D卷积：卷积核3*3*3，步长1，填充1
            torch.nn.BatchNorm3d(out_channels), # 批归一化
            torch.nn.ReLU() # ReLU激活函数
        )

    def forward(self, x):
        out = self.conv20(x)
        out = self.conv21(out)
        out = self.conv22(out)
        return out
